#include "DTIFiberTrackingCPU.h"

#include <string>
#include <omp.h>

#ifndef THREADS_COUNT
#define THREADS_COUNT 16
#endif

void CDTIFiberTrackingCPU::FiberTracking_TEND(const char *fiberPathname, const uint scale)
{
	int sw = m_width * scale;
	int sh = m_height * scale;
	int sd = m_depth * scale;

	FILE *fp = fopen(fiberPathname, "w");
	fprintf(fp, "%d %d %d %d %f %f\n", sw, sh, sd, scale, m_stepSize, m_maxLength);

	std::string pathname = fiberPathname;
	size_t pos = pathname.find_last_of(SLASH);
	std::string path = pathname.substr(0, pos+1);

	for (int z = 0; z < sd; ++z) {
		CFiber *fibers = new CFiber[sw*sh];
		Seed *seeds = new Seed[sw*sh];
		GenerateSeeds_TEND(seeds, scale, z);

		for (int y = 0; y < sh; ++y) {
			omp_set_num_threads(THREADS_COUNT);
#pragma omp parallel for
			for (int x = 0; x < sw; ++x) {
				int index = y * sw + x;

				TrackingOneFiber_TEND(&(seeds[index]), &(fibers[index]));
			}
		}

		char confPathname[PATHNAME_SIZE];
		char vertPathname[PATHNAME_SIZE];
		sprintf(confPathname, "%sslice_z=%08.04f.conf", path.c_str(), (float)z/(float)scale);
		sprintf(vertPathname, "%sslice_z=%08.04f.vert", path.c_str(), (float)z/(float)scale);
		fprintf(fp, "%s\t%s\n", confPathname, vertPathname);
		
		SaveFibers(fibers, sw, sh, confPathname, vertPathname);

		SafeDeleteArray(seeds);
		SafeDeleteArray(fibers);
	}

	fclose(fp);
}


void CDTIFiberTrackingCPU::GenerateSeeds_TEND(Seed *seeds, const uint scale, const int z)
{
	int sw = m_width * scale;
	int sh = m_height * scale;
	int sd = m_depth * scale;

	omp_set_num_threads(16);
#pragma omp parallel for
	for (int y = 0; y < sh; ++y) {
		for (int x = 0; x < sw; ++x) {
			int index = y * sw + x;

			CPoint3F pos;
			pos.m_x = (float)x/(float)scale - 0.5f + 1.0f / ((float)scale * 2.0f);
			pos.m_y = (float)y/(float)scale - 0.5f + 1.0f / ((float)scale * 2.0f);
			pos.m_z = (float)z/(float)scale - 0.5f + 1.0f / ((float)scale * 2.0f);

			if (m_pMask->GetScalar(&pos) < 0.5f) {
				seeds[index].pos = pos;
				seeds[index].dir = CVector3F(0.0f, 0.0f, 0.0f);
			} else {
				float dti[7];
				uint xx = uint(pos.m_x + 0.5f);
				uint yy = uint(pos.m_y + 0.5f);
				uint zz = uint(pos.m_z + 0.5f);
				m_pDTI->GetDTI(xx, yy, zz, dti);
				CVector3F dir;
				m_pDTI->DTI2Direction(dti, &dir);

				seeds[index].pos = pos;
				seeds[index].dir = dir;
			}
		}
	}
}

void CDTIFiberTrackingCPU::TrackingOneFiber_TEND(Seed *seed, CFiber *fiber)
{
	float p = fabs(seed->dir.m_x) + fabs(seed->dir.m_y) + fabs(seed->dir.m_z);
	if (p < 0.5f) {
		fiber->CreateFiber(&(seed->pos), &(seed->pos), 1, 1, seed);
		return;
	}


	uint maxSteps = (uint)(m_maxLength / (m_stepSize * 10.0f));
	uint fnum, bnum;
	CPoint3F *pF = new CPoint3F[maxSteps];
	CPoint3F *pB = new CPoint3F[maxSteps];
	CPoint3F lastPos;
	CVector3F lastDir;
	int lastIndex;

	/* tracking forward */
	lastPos = seed->pos;
	lastDir = seed->dir;
	lastIndex = ComputeIndex(&lastPos);
	pF[0] = lastPos;
	for (fnum = 1; fnum < maxSteps; ++fnum) {
		CPoint3F nextPos;
		CVector3F nextDir;
		int nextIndex;

		bool flag = true;
		for (int step = 0; step < 10; ++step) {
			nextPos.m_x = lastPos.m_x + lastDir.m_x * m_stepSize;
			nextPos.m_y = lastPos.m_y + lastDir.m_y * m_stepSize;
			nextPos.m_z = lastPos.m_z + lastDir.m_z * m_stepSize;

			nextIndex = ComputeIndex(&nextPos);
			if (nextIndex == lastIndex) {
				if (m_pFA->GetScalar(&nextPos) < m_faThreshold) {
					flag = false;
					break;
				}
				nextDir = lastDir;
			} else {
				if (GetNextDirection_TEND(&nextPos, &lastDir, &nextDir) == 0) {
					flag = false;
					break;
				}
			}

			lastDir = nextDir;
			lastPos = nextPos;
			lastIndex = nextIndex;
		}
		
		if (flag == false)
			break;

		pF[fnum] = nextPos;
		lastDir = nextDir;
		lastPos = nextPos;
		lastIndex = nextIndex;
	}


	/* tracking backward */
	lastPos = seed->pos;
	lastDir = seed->dir;
	lastIndex = ComputeIndex(&lastPos);
	pB[0] = lastPos;
	for (bnum = 1; bnum < maxSteps; ++bnum) {
		CPoint3F nextPos;
		CVector3F nextDir;
		int nextIndex;

		bool flag = true;
		for (int step = 0; step < 10; ++step) {
			nextPos.m_x = lastPos.m_x - lastDir.m_x * m_stepSize;
			nextPos.m_y = lastPos.m_y - lastDir.m_y * m_stepSize;
			nextPos.m_z = lastPos.m_z - lastDir.m_z * m_stepSize;

			nextIndex = ComputeIndex(&nextPos);
			if (nextIndex == lastIndex) {
				if (m_pFA->GetScalar(&nextPos) < m_faThreshold) {
					flag = false;
					break;
				}
				nextDir = lastDir;
			} else {
				if (GetNextDirection_TEND(&nextPos, &lastDir, &nextDir) == 0) {
					flag = false;
					break;
				}
			}

			lastDir = nextDir;
			lastPos = nextPos;
			lastIndex = nextIndex;
		}
		
		if (flag == false)
			break;
		
		pB[bnum] = nextPos;
		lastDir = nextDir;
		lastPos = nextPos;
	}

	fiber->CreateFiber(pF, pB, fnum, bnum, seed);

	
	SafeDeleteArray(pF);
	SafeDeleteArray(pB);
}

int CDTIFiberTrackingCPU::GetNextDirection_TEND(CPoint3F *pos, CVector3F *in, CVector3F *out)
{
	if (m_pMask->GetScalar(pos) < 0.5f) {
		return 0;
	}

	if (m_pFA->GetScalar(pos) < m_faThreshold) {
		return 0;
	}

	float dti[7];
	uint xx = uint(pos->m_x + 0.5f);
	uint yy = uint(pos->m_y + 0.5f);
	uint zz = uint(pos->m_z + 0.5f);
	m_pDTI->GetDTI(xx, yy, zz, dti);

	CVector3F dir;
	m_pDTI->DTI2Direction(dti, &dir);

	CVector3F temp;
	temp.m_x = dti[1] * in->m_x + dti[2] * in->m_y + dti[3] * in->m_z;
	temp.m_y = dti[2] * in->m_x + dti[4] * in->m_y + dti[5] * in->m_z;
	temp.m_z = dti[3] * in->m_x + dti[5] * in->m_y + dti[6] * in->m_z;

	out->m_x = m_tendF * dir.m_x + (1.0f - m_tendF) * ((1.0f - m_tendG) * in->m_x + m_tendG * temp.m_x);
	out->m_y = m_tendF * dir.m_y + (1.0f - m_tendF) * ((1.0f - m_tendG) * in->m_y + m_tendG * temp.m_y);
	out->m_z = m_tendF * dir.m_z + (1.0f - m_tendF) * ((1.0f - m_tendG) * in->m_z + m_tendG * temp.m_z);

	Normalize(out);

	float v = InnerProduct(in, out);
	if (fabs(v) < m_angleThreshold) {
		return 0;
	}

	if (v < 0.0f) {
		out->m_x = -out->m_x;
		out->m_y = -out->m_y;
		out->m_z = -out->m_z;
	}

	return 1;
}
